import streamlit as st

# 设置页面配置（必须是第一个st命令）
st.set_page_config(
    page_title="医用试管智能识别与监管系统",
    page_icon="🧪",
    layout="wide"
)

import cv2
import numpy as np
from PIL import Image
import os
import datetime
import traceback
import warnings
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 设置环境变量以提高兼容性
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"  # 启用MPS回退
os.environ["TORCH_SHOW_CPP_STACKTRACES"] = "0"   # 禁用C++堆栈跟踪
os.environ["TORCH_WEIGHTS_ONLY"] = "0"           # 关闭weights_only限制

# 抑制不必要的警告
warnings.filterwarnings("ignore", message=".*detected in torch._classes.*")
warnings.filterwarnings("ignore", message=".*gpu_mps.*")
warnings.filterwarnings("ignore", message=".*does not exist.*")
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)

# 设置环境变量以改善CUDA兼容性
os.environ["TORCH_CUDA_ARCH_LIST"] = "3.5;5.0;6.0;7.0;7.5;8.0;8.6"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"

# 导入自定义模块
try:
    from utils.image_processing import process_image, detect_tubes
    st.info("成功导入图像处理模块")
except ImportError as e:
    st.error(f"导入图像处理模块失败: {str(e)}")
    def process_image(image): return image
    def detect_tubes(image, model): return image, []

try:
    from database.db_manager import init_db, save_tube_data, get_all_records
    st.info("成功导入数据库模块")
except ImportError as e:
    st.error(f"导入数据库模块失败: {str(e)}")
    def init_db(): return False
    def save_tube_data(*args, **kwargs): return None
    def get_all_records(): return []

try:
    from models.model_handler import load_model
    st.info("成功导入模型处理模块")
except ImportError as e:
    st.error(f"导入模型处理模块失败: {str(e)}")
    def load_model(): return None

# 初始化数据库
db_initialized = init_db()

# 配置页面
st.title("医用试管智能识别与监管系统")
st.markdown("""
本系统使用计算机视觉技术识别医用试管，可以检测:
- 试管数量和位置
- 试管管冒颜色
- 液体位置
- 试管标签内容
""")

# 加载模型
try:
    # 禁止输出调试信息
    logging.getLogger("ultralytics").setLevel(logging.ERROR)
    
    # 使用上下文管理器捕获和隐藏模型加载过程中的警告和错误输出
    with st.spinner('正在加载多功能模型...'):
        # 加载多个模型：检测、分割和分类
        models = load_model()
        if not models:
            st.warning("模型加载失败，将使用传统方法识别试管")
        else:
            st.success("模型加载成功")
except Exception as e:
    st.error(f"模型加载过程中发生错误: {str(e)}")
    if "no attribute '__module__'" in str(e):
        st.error("PyTorch模型加载错误: 可能是PyTorch和ultralytics版本兼容性问题")
        st.info("系统将使用传统方法进行试管识别")
    # 确保models变量已定义，使用传统方法模型
    from models.model_handler import TraditionalModel
    models = {'main': TraditionalModel()}

# 边栏
with st.sidebar:
    st.title("操作面板")
    st.markdown("---")
    
    # 选择输入方式
    input_method = st.radio("选择输入方式:", ["上传图片", "使用摄像头"])
    
    if input_method == "上传图片":
        # 上传图片
        uploaded_file = st.file_uploader("上传试管图片", type=["jpg", "jpeg", "png"])
        process_btn = st.button("开始识别")
    else:
        # 摄像头设置提示
        st.info("请使用下方摄像头拍摄试管照片，拍照后将自动进行识别")
    
    st.markdown("---")
    st.write("系统操作：")
    view_db_btn = st.button("查看数据库记录", disabled=not db_initialized)

# 主页面
st.header("医用试管智能识别系统")

# 状态信息
if not db_initialized:
    st.warning("数据库连接失败，系统将在离线模式下运行，识别结果将不会被保存")

# 创建两列布局
col1, col2 = st.columns(2)

# 定义变量以避免未定义错误
image = None
save_path = None
camera_image = None

# 处理摄像头输入
if input_method == "使用摄像头":
    with col1:
        st.subheader("摄像头捕获")
        # 使用Streamlit的camera_input组件
        camera_image = st.camera_input("拍摄试管照片")
        
        if camera_image is not None:
            try:
                # 显示拍摄的照片
                image = Image.open(camera_image)
                # 保存拍摄的图片
                os.makedirs(os.path.join("data", "uploads"), exist_ok=True)
                timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                save_path = os.path.join("data", "uploads", f"camera_{timestamp}.jpg")
                with open(save_path, "wb") as f:
                    f.write(camera_image.getbuffer())
                
                st.success("照片已拍摄并保存，将进行识别")
            except Exception as e:
                st.error(f"摄像头图像处理错误: {str(e)}")
                image = None

# 处理上传图片
if input_method == "上传图片":
    with col1:
        st.subheader("原始图像")
        if uploaded_file is not None:
            try:
                # 保存上传的图片
                os.makedirs(os.path.join("data", "uploads"), exist_ok=True)
                save_path = os.path.join("data", "uploads", uploaded_file.name)
                with open(save_path, "wb") as f:
                    f.write(uploaded_file.getbuffer())
                
                # 显示原图
                image = Image.open(uploaded_file)
                st.image(image, caption="上传的图片", width=None)
            except Exception as e:
                st.error(f"图像处理错误: {str(e)}")
                image = None

# 处理试管识别
with col2:
    st.subheader("识别结果")
    # 判断是否需要处理图像
    should_process = False
    
    if input_method == "上传图片" and uploaded_file is not None and process_btn and image is not None:
        should_process = True
    elif input_method == "使用摄像头" and camera_image is not None and image is not None:
        should_process = True
    
    if should_process:
        with st.spinner('正在处理图像...'):
            try:
                # 处理图像
                image_array = np.array(image)
                processed_image, tube_data = detect_tubes(image_array, models)
                
                # 保存处理后的图像
                os.makedirs(os.path.join("data", "processed"), exist_ok=True)
                if input_method == "上传图片":
                    processed_path = os.path.join("data", "processed", f"proc_{uploaded_file.name}")
                else:
                    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                    processed_path = os.path.join("data", "processed", f"proc_camera_{timestamp}.jpg")
                
                cv2.imwrite(processed_path, cv2.cvtColor(processed_image, cv2.COLOR_RGB2BGR))
                
                # 显示处理后的图像
                st.image(processed_image, caption="识别结果", width=None)
                
                # 显示详细信息
                if tube_data and len(tube_data) > 0:
                    # 创建一个表格显示详细信息
                    tube_info = []
                    for tube in tube_data:
                        tube_info.append({
                            "ID": tube.get('id', ''),
                            "试管长度": f"{tube.get('length', 0)}像素",
                            "管冒颜色": tube.get('cap_color', '未知'),
                            "液体位置": tube.get('liquid_position', '未知'),
                            "试管标签": tube.get('label_text', '无法识别')
                        })
                    
                    st.write("试管详细信息:")
                    st.table(tube_info)
                
                # 保存到数据库
                if tube_data and save_path is not None:
                    success_count = 0
                    # 限制保存的试管数量，防止过多重复记录
                    max_tubes = min(7, len(tube_data))
                    
                    # 只保存前max_tubes个试管数据
                    for i, tube in enumerate(tube_data[:max_tubes]):
                        # 使用正确的字段名
                        record_id = save_tube_data(
                            image_path=save_path,
                            tube_length=tube.get('length', 0),
                            cap_color=tube.get('cap_color', '未知'),
                            liquid_position=tube.get('liquid_position', '未知'),
                            label_text=tube.get('label_text', ''),
                            confidence=tube.get('confidence', 0.5)
                        )
                        if record_id:
                            success_count += 1
                    
                    if db_initialized:
                        st.success(f"成功识别 {success_count} 个试管并存入数据库")
                    else:
                        st.info(f"成功识别 {success_count} 个试管，但由于数据库未连接，数据未保存")
                else:
                    st.warning("未检测到任何试管")
            except Exception as e:
                st.error(f"图像处理过程中出错: {str(e)}")
                st.error(traceback.format_exc())

# 显示数据库记录
if view_db_btn and db_initialized:
    st.markdown("---")
    st.subheader("数据库记录")
    
    records = get_all_records()
    if records:
        # 创建数据表格
        data = []
        for record in records:
            data.append({
                "ID": record.id,
                "图像路径": record.image_path,
                "试管长度(mm)": record.tube_length,
                "试管冒颜色": record.cap_color,
                "液体位置": record.liquid_position,
                "标签内容": record.label_text,
                "置信度": f"{record.confidence:.2f}",
                "记录时间": record.timestamp
            })
        
        st.dataframe(data)
    else:
        st.info("数据库中暂无记录") 